import time
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from model_interfaces import IASRModel
import torch

class FunASRModel(IASRModel):
    def __init__(self, device: torch.device):
        print("正在初始化FunASR模型...")
        start_time = time.time()
        # 使用针对中文优化的Paraformer-large模型
        self.model = pipeline(
            task=Tasks.auto_speech_recognition,
            model='damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch',
            device=device.type # 'cuda' or 'cpu'
        )
        end_time = time.time()
        print(f"FunASR模型加载完毕。耗时: {end_time - start_time:.2f} 秒")
        self._transcript = ""
        self._word_locations = []

    def processAudio(self, audio: torch.Tensor, sample_rate: int = 16000) -> None:
        """
        处理音频张量以生成转录和带时间戳的词语信息。
        """
        if audio.ndim > 1:
            audio = audio.squeeze(0) # 转换为一维

        audio_data = audio.cpu().numpy()

        try:
            # FunASR pipeline直接接受numpy数组
            result = self.model(audio_data)
            
            # 处理结果，FunASR可能返回单个结果或结果列表
            if isinstance(result, list) and len(result) > 0:
                result = result[0]
            
            self._transcript = result.get('text', '') if isinstance(result, dict) else str(result)
            
            # 尝试获取时间戳信息
            if isinstance(result, dict):
                raw_word_info = result.get('timestamp', []) or result.get('word_info', [])
                self._word_locations = self._process_timestamps(raw_word_info, sample_rate)
            else:
                self._word_locations = []
                
            print(f"ASR转录结果: {self._transcript}")

        except Exception as e:
            print(f"FunASR处理音频时发生错误: {e}")
            print(f"错误类型: {type(e).__name__}")
            import traceback
            traceback.print_exc()
            self._transcript = ""
            self._word_locations = []

    def getTranscript(self) -> str:
        return self._transcript

    def getWordLocations(self) -> list:
        """
        返回处理后的时间戳信息，格式为 [{'word': str, 'start_ts': int, 'end_ts': int}, ...]
        """
        return self._word_locations

    def _process_timestamps(self, word_info: list, sample_rate: int) -> list:
        """
        将FunASR的时间戳转换为采样点数。
        支持多种可能的时间戳格式。
        """
        processed_locations = []
        
        if not word_info:
            return processed_locations
            
        for item in word_info:
            try:
                if isinstance(item, dict):
                    # 方式1: {'word': '今天', 'start': 100, 'end': 300}
                    word = item.get('word', item.get('text', ''))
                    start = item.get('start', item.get('start_time', 0))
                    end = item.get('end', item.get('end_time', 0))
                elif isinstance(item, (list, tuple)) and len(item) >= 3:
                    # 方式2: [word, start, end]
                    word, start, end = item[0], item[1], item[2]
                else:
                    continue
                    
                # 时间戳可能是毫秒或秒，根据数值大小判断
                if start > 10000:  # 假设是毫秒
                    start_ts = int((start / 1000) * sample_rate)
                    end_ts = int((end / 1000) * sample_rate)
                else:  # 假设是秒
                    start_ts = int(start * sample_rate)
                    end_ts = int(end * sample_rate)
                    
                processed_locations.append({
                    'word': str(word),
                    'start_ts': start_ts,
                    'end_ts': end_ts
                })
            except Exception as e:
                print(f"处理时间戳时发生错误: {e}, item: {item}")
                continue
                
        return processed_locations

def get_asr_model(device: torch.device):
    """
    模型工厂函数
    """
    return FunASRModel(device=device) 